# -*- coding:utf-8 -*-
import asyncio
import csv
import json

# @Time    : 2023/5/13 02:24
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : user_info_extract.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import traceback

from bot.insurance_planner_gpt.agent.base_agent import LLMAgent
from common.log import logger
import os
import pandas as pd
import docx
import textract

curPath = os.path.abspath(os.path.dirname(__file__))

default_template = """
帮我按照知识库的形式把下面的产品文档内容整理成问答格式,Q和A分别代表问题和答案，你需要一步一步仔细思考，把产品文档仔细阅读完后基于语义给出答案，内容可以多一点但不要漏掉，也不要被举例影响。
1.提出的问题基于这些问题发散：这个产品都包含什么保险责任？每一项保险责任具体内容是什么？如果不能直接获取答案需要你仔细思考然后总结
其中'==='之间的是产品文档。
输出结果是json，格式是:{{"Q":"","A":""}} 文档没有的话不返回

举例：
---
名称：i无忧2.0重疾险
保障责任：（可选）疾病关爱保险金
一、轻症疾病关爱保险金
被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的轻症疾病（一种或多种），且此前未确诊患有本合同约定的重大疾病（一种或多种），则保险公司在给付轻症疾病保险金的同时，按基本保险金额的10%给付轻症疾病关爱保险金。轻症疾病关爱保险金给付以一次为限。
二、中症疾病关爱保险金
被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的中症疾病（一种或多种），且此前未确诊患有本合同约定的重大疾病（一种或多种），则保险公司在给付中症疾病保险金的同时，按基本保险金额的30%给付中症疾病关爱保险金。中症疾病关爱保险金给付以一次为限。
三、重大疾病关爱保险金
被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的重大疾病（一种或多种），则保险公司在给付第一次重大疾病保险金的同时，按基本保险金额的80%给付重大疾病关爱保险金。重大疾病关爱保险金给付以一次为限。

输出：
{{"Q":"i无忧2.0重疾险都包含什么服务？","A":"疾病关爱保险金 其中包含轻症疾病关爱保险金，中症疾病关爱保险金，重大疾病关爱保险金"}}
{{"Q":"什么是轻症疾病关爱保险金？","A":"被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的轻症疾病（一种或多种），且此前未确诊患有本合同约定的重大疾病（一种或多种），则保险公司在给付轻症疾病保险金的同时，按基本保险金额的10%给付轻症疾病关爱保险金。轻症疾病关爱保险金给付以一次为限"}}
{{"Q":"什么是中症疾病关爱保险金？","A":"被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的中症疾病（一种或多种），且此前未确诊患有本合同约定的重大疾病（一种或多种），则保险公司在给付中症疾病保险金的同时，按基本保险金额的30%给付中症疾病关爱保险金。中症疾病关爱保险金给付以一次为限。"}}
{{"Q":"什么是重大疾病关爱保险金？","A":"被保险人于等待期90天后，且于年满60周岁后的首个保单年生效对应日零时之前，经保险公司认可的医院确诊初次患有本合同约定的重大疾病（一种或多种），则保险公司在给付第一次重大疾病保险金的同时，按基本保险金额的80%给付重大疾病关爱保险金。重大疾病关爱保险金给付以一次为限。"}}
{{"Q":"轻症疾病关爱保险金能赔付几次？","A":"轻症疾病关爱保险金给付以一次为限"}}
举例结束（任务忽略）；
---
产品文档：
===
{document}
===
从这些文档里提取到的内容为：
        """


class Document2Knowledge(LLMAgent):

    def __init__(self, document):
        self.prompt = default_template
        self.prompt = self.prompt.format(document=document)
        super().__init__(self.prompt)

    def parse_knowledge(self, knowledge: str, file_name, saveAsHistory=True, ):
        q_and_a_list = []
        try:
            lines = knowledge.split('\n')
            for line in lines:
                if line:
                    data = json.loads(line)
                    data['product'] = file_name.split('.t')[0]
                    q_and_a_list.append(data)

            if saveAsHistory and q_and_a_list:
                with open("q_and_a_history.json", "a") as f:
                    json.dump(q_and_a_list, f, ensure_ascii=False)
                    f.write("\n")
        except Exception as ee:
            pass

        return q_and_a_list


def write_knowledge_to_csv(targetKnowledgePath: str, knowledges: [], writeFile="保险产品说明书.csv"):
    try:

        loadCount = {}
        for kDict in knowledges:
            product = kDict.get('product')
            q = kDict.get('Q')
            a = kDict.get('A')

            content = [product, q, a]

            if os.path.exists(os.path.join(targetKnowledgePath, writeFile)) is False:
                with open(os.path.join(targetKnowledgePath, writeFile), "a", newline="") as f:
                    w = csv.writer(f, delimiter=",")
                    w.writerow(['产品名称', '产品问题', '知识内容'])
                f.close()

            with open(os.path.join(targetKnowledgePath, writeFile), "a", newline="") as f:
                w = csv.writer(f, delimiter=",")
                w.writerow(content)
            f.close()
            if loadCount.get(type):
                loadCount[type] = int(loadCount.get(type)) + 1
            else:
                loadCount[type] = 1
        logger.info('写入知识库 {} 知识'.format(loadCount))
    except Exception as ee:
        logger.error(traceback.format_exc())


if __name__ == '__main__':

    need_append = False

    if need_append:
        pass
        # os.remove(curPath + '/raw_data/product_items/保险产品说明书.csv')

    # 素材库提取
    file_path = curPath + '/raw_data/products_doc/'

    if os.path.isdir(file_path):
        result_list = []
        for file in os.listdir(file_path):
            if not file.endswith('.txt'):
                continue
            with open(file_path + file, 'r', encoding='utf-8') as f:
                content = f.read()

            document_knowledge = Document2Knowledge(content)

            result = asyncio.run(document_knowledge.achat_auto_llm(type='gpt'))
            result = result.replace('```json', '').replace('```', '')
            print(result)
            if result:
                knowledge = document_knowledge.parse_knowledge(result, file)
                result_list.extend(knowledge)
            # 写知识库csv文件
        write_knowledge_to_csv(curPath + '/raw_data/products_doc_csv/', result_list)
